import random

import numpy as np
import torch.cuda
from numpy.random import shuffle
from torch import nn

from common_models import gan
from common_models.MLP import MLP
from common_utils.utils import l2_normalize
import torch.nn.functional as F

device = 'cuda' if torch.cuda.is_available() else 'cpu'


class ALEModel(nn.Module):
    def __init__(self, args, feat_dim, attr_dim):
        super(ALEModel, self).__init__()
        # 生成器
        self.netG = gan.MLP_Generator(input_dim=args.attSize, output_dim=args.resSize,
                                 layers=[args.ngh, args.ngh * 2, args.ngh]).to(device)
        # 输入维度为图片特征维度，输出维度为属性的数量
        self.label_classifier = MLP(feat_dim, 50, relu=True, dropout=args.dropout, norm=args.norm).to(device)

    def forward(self, imgs, attrs=None,labels=None, sig=None):
        if imgs is None:
            fake_imgs = self.netG(attrs)
        # 根据图片特征学习对象的属性
        preds = self.label_classifier(imgs)  # 属性评分
        # loss_fn = torch.nn.MSELoss()# 训练速度快，训练集过拟合
        loss_fn = nn.CrossEntropyLoss()
        loss = loss_fn(preds, labels)
        return loss.mean(), preds
